'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 2_condition_and_loop_graph.py
* @Time: 2025/11/26
* @All Rights Reserve By Brtc
'''
import json
from typing import TypedDict, List, Annotated

import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSearchAPIWrapper, GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.messages import ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.constants import START
from langgraph.graph import add_messages, StateGraph
from pydantic import BaseModel, Field
from typing_extensions import Literal

dotenv.load_dotenv()
#工具绑定
class GoogleSearchSchema(BaseModel):
    query:str = Field(description="谷歌搜索的参数")
class DalleArgsSchema(BaseModel):
    query:str = Field(description="图片生成指令")
google_search_tool = GoogleSerperRun(
    name = "google_search_tool",
    description="一个低成本的谷歌搜索api",
    api_wrapper=GoogleSerperAPIWrapper(),
    args_schema=GoogleSearchSchema
)
dalle_tool = OpenAIDALLEImageGenerationTool(
    name = "dalle_tool",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    args_schema=DalleArgsSchema
)

bot_tools = [google_search_tool, dalle_tool]
llm = ChatOpenAI(model="gpt-4o-mini")
llm_with_tools = llm.bind_tools(bot_tools)

# 图架构的状态定义
class  ChatBotState(TypedDict):
    messages: Annotated[list, add_messages]

def chat_bot(state: ChatBotState)->ChatBotState:
    """大模型节点"""
    ai_message = llm_with_tools.invoke(state["messages"])
    return {"messages":[ai_message]}

def tool_exe(state:ChatBotState)->ChatBotState:
    """工具执行节点"""
    tool_by_name = {tool.name:tool for tool in bot_tools}
    tool_calls = state["messages"][-1].tool_calls # 这个消息就是AI传递的工具调用消息

    messages = []
    for tool_call in tool_calls:
        tool = tool_by_name[tool_call["name"]]# 获取调用工具实例
        #构建工具调用消息
        messages.append(ToolMessage(
            tool_call_id = tool_call["id"],
            content = json.dumps(tool.invoke(tool_call["args"])),
            name = tool_call["name"],
        ))

    return {"messages":messages}

def route(state:ChatBotState)->Literal["tool_exe", "__end__"]:
    """动态选择工具或者结束边"""
    ai_message = state["messages"][-1] # state 的最后一条消息，AI消息
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tool_exe"
    return "__end__"



"""图架构  
                <------loop-- tool_exe
                |              |
START -------> llm --------> route(condition_edge) ------->end
"""

#1、创建图架构
graph_builder = StateGraph(ChatBotState)
#2、添加节点或者边
graph_builder.add_node("llm", chat_bot)
graph_builder.add_node("tool_exe", tool_exe)
#3、添加边
graph_builder.add_edge(START, "llm")
graph_builder.add_conditional_edges("llm", route)
graph_builder.add_edge("tool_exe", "llm")
#4、编译图架构
graph_app = graph_builder.compile()
#5、运行图架构
"""  
输入:state:
{
"messages":[("human","帮我画一幅 老人与海的图片")]
}

输出:
{
"messages":[("human","帮我画一幅 老人与海的图片"), ai , too_message, ai]
}

"""
state = graph_app.invoke({"messages":[("human","帮我画一幅 老人与海的图片")]})

for message in state["messages"]:
    print(f"消息类型:{message.type}")
    if hasattr(message,"tool_calls") and len(message.tool_calls) > 0:
        print(f"工具调用参数:{message.tool_calls}")
    print(f"消息内容:{message.content}")
    print("=================================================")